#!/bin/python
# -*-coding: utf-8-*-

"""
Downscaling meterological field
"""
import sys
sys.path.append("..")

import utils.raster as raster
import utils.common as common
import rfe
import wind_vector

import os
import math
import datetime

import numpy as np
import gdal
import matplotlib.pyplot as plt
import cv2

def downscale_metero_filed(root,out_root,start_date,end_date):
	"""
	Args:
		root
		start_date: string of date, i.e. 2013-01-01
		end_date: string of date, i.e. 2013-01-01
	Returns:

	Raises:
		RuntimeError: 
	"""
	interoplate_method = "bilinear"
	temp_lapse = {1:4.4,2:5.9,3:7.1,4:7.8,5:8.1,6:8.2,7:8.1,8:8.1,9:7.7,10:6.8,11:5.5,12:4.7}
	precp_lapse = {1:0.35,2:0.35,3:0.35,12:0.35,4:0.3,11:0.3,5:0.25,10:0.25,6:0.2,7:0.2,8:0.2,9:0.2}
	lambda_lapse = {1:0.41,2:0.42,3:0.4,11:0.4,12:0.4,4:0.39,5:0.38,6:0.36,9:0.36,7:0.33,8:0.33,10:0.37}
	# for relative humidity
	a = 611.21
	b = 17.502
	c = 240.97
	# for wind speed & direction
	# gamma_c plus gamma_s should be equal 1
	gamma_c = 0.5
	gamma_s = 0.5
	## DEM and topographical factors
	dem = gdal.Open(os.path.join(root,"dem","dem.tif"))
	dem_info = raster.get_raster_info(dem)
	dem_proc_scale = round(math.cos((dem_info[1][3] + dem.RasterXSize / 2 * dem_info[1][4] + dem.RasterYSize / 2 * dem_info[1][5]) * math.pi / 180.0) * 111120)
	dem_aspect = gdal.DEMProcessing("",dem,"aspect",scale=dem_proc_scale,format="MEM")
	dem_slope = gdal.DEMProcessing("",dem,"slope",scale=dem_proc_scale,format="MEM")

	geopotential = gdal.Open(os.path.join(root,"geopotential","geopotential.tif"))
	geop_inter = raster.resample(geopotential,dem,interoplate_method)

	wind_curvature = __calc_curvature__(dem,dem_proc_scale)

	elev_diff = np.ma.array(dem.GetRasterBand(1).ReadAsArray() - geop_inter.GetRasterBand(1).ReadAsArray(),mask = dem.GetRasterBand(1).ReadAsArray() == dem_info[0],fill_value = 0)

	sdate = datetime.datetime.strptime(start_date,"%Y-%m-%d")
	edate = datetime.datetime.strptime(end_date,"%Y-%m-%d")

	for date_delta in range(0,(edate-sdate).days+1):
		cur_date = sdate + datetime.timedelta(days = date_delta)
		file_posix = datetime.datetime.strftime(cur_date,"%Y%j")
		## Temperature
		temp_era = gdal.Open(os.path.join(root,"t2m","t2m_%s.tif"%file_posix))

		temp_inter = raster.resample(temp_era,dem,interoplate_method)

		temp_dem = temp_inter.GetRasterBand(1).ReadAsArray() - 273.15 - (elev_diff / 1000.0) * temp_lapse[cur_date.month]

		out_temp_path = os.path.join(out_root,"t2m")

		common.is_exists(out_temp_path)

		raster.save(os.path.join(out_temp_path,"t2m_%s.tif"%file_posix),temp_dem,-999.0,dem_info[1],dem_info[2])

		# Max Temperature
		temp_era = gdal.Open(os.path.join(root,"m2m","m2m_%s.tif"%file_posix))

		temp_inter = raster.resample(temp_era,dem,interoplate_method)

		temp_dem = temp_inter.GetRasterBand(1).ReadAsArray() - 273.15 - (elev_diff / 1000.0) * temp_lapse[cur_date.month]

		out_temp_path = os.path.join(out_root,"m2m")

		common.is_exists(out_temp_path)

		raster.save(os.path.join(out_temp_path,"m2m_%s.tif"%file_posix),temp_dem,-999.0,dem_info[1],dem_info[2])

		## Relative Humidity
		dew_temp = gdal.Open(os.path.join(root,"d2m","d2m_%s.tif"%file_posix))
		dew_temp_inter = raster.resample(dew_temp,dem,interoplate_method)

		dew_temp_dem = dew_temp_inter.GetRasterBand(1).ReadAsArray() - 273.15 - (elev_diff / 1000.0) * lambda_lapse[cur_date.month] * (c / b)

		rh = 100 * (np.exp(((b * dew_temp_dem)/(c + dew_temp_dem))) / np.exp(((b * temp_dem) / (c + temp_dem))))
		# rh = dew_temp_dem / temp_dem

		out_rh_path = os.path.join(out_root,"rh")

		common.is_exists(out_rh_path)
		# TODO:
		# the rh value will greater than 100%
		# raster.save(os.path.join(out_rh_path,"d2m_%s.tif"%file_posix),dew_temp_dem,-999.0,dem_info[1],dem_info[2])
		rh[rh > 100.0] = 100.0

		raster.save(os.path.join(out_rh_path,"rh_%s.tif"%file_posix),rh,-999.0,dem_info[1],dem_info[2])

		## Wind Speed
		u = gdal.Open(os.path.join(root,"u10","u10_%s.tif"%file_posix))
		v = gdal.Open(os.path.join(root,"v10", "v10_%s.tif"%file_posix))

		u_inter = raster.resample(u,dem,interoplate_method)
		v_inter = raster.resample(v,dem,interoplate_method)

		(wns,wnd) = wind_vector.wind_vector(u_inter.GetRasterBand(1).ReadAsArray(),v_inter.GetRasterBand(1).ReadAsArray())

		wind_slope = dem_slope.GetRasterBand(1).ReadAsArray() * np.cos((wnd - dem_aspect.GetRasterBand(1).ReadAsArray())* np.pi / 180.0)

		omega_s = (wind_slope - (np.max(wind_slope)+np.min(wind_slope))/2.0)/(np.max(wind_slope) - np.min(wind_slope))

		w_s = (1 + gamma_c * wind_curvature + gamma_s * omega_s)

		wns_dem = wns * (1 + gamma_c * wind_curvature + gamma_s * omega_s)

		wnd_dem = wnd - 0.5 * omega_s * np.sin((2*(dem_aspect.GetRasterBand(1).ReadAsArray() - wnd)) * np.pi / 180.)

		out_wns_path = os.path.join(out_root,"wns")
		common.is_exists(out_wns_path)

		out_wnd_path = os.path.join(out_root,"wnd")
		common.is_exists(out_wnd_path)

		raster.save(os.path.join(out_wns_path,"wns_%s.tif"%file_posix),wns_dem,-999.0,dem_info[1],dem_info[2])
		raster.save(os.path.join(out_wnd_path,"wnd_%s.tif"%file_posix),wnd_dem,-999.0,dem_info[1],dem_info[2])

		## Precipitation
		rfe_data = rfe.parse(os.path.join(root,"rfe","cpc_rfe_v2.0_sa_dly.bin.%s.gz"%datetime.datetime.strftime(cur_date,"%Y%m%d")))
		rfe_interp = raster.resample(rfe_data,dem,interoplate_method)
		rfe_info = raster.get_raster_info(rfe_data)

		rfe_dem = rfe_interp.GetRasterBand(1).ReadAsArray() * ((1 + precp_lapse[cur_date.month] * (elev_diff / 1000.0))/ (1 - precp_lapse[cur_date.month] * (elev_diff / 1000.0)))

		out_rfe_path = os.path.join(out_root,"rfe")
		common.is_exists(out_rfe_path)
		raster.save(os.path.join(out_rfe_path,"rfe_%s.tif"%file_posix),rfe_dem,-999.0,dem_info[1],dem_info[2])
		# raster.save(os.path.join(out_rfe_path,"rfei_%s.tif"%file_posix),rfe_data.GetRasterBand(1).ReadAsArray(),-999.0,rfe_info[1],rfe_info[2])
		# raster.save(os.path.join(out_rfe_path,"rfeii_%s.tif"%file_posix),rfe_interp.GetRasterBand(1).ReadAsArray(),-999.0,dem_info[1],dem_info[2])

def __calc_curvature__(dem,eita):
	geo_trans = dem.GetGeoTransform()
	dem_array = dem.GetRasterBand(1).ReadAsArray().astype(float)
	fill_value = dem.GetRasterBand(1).GetNoDataValue()
	dem_array[dem_array == fill_value] = 0

	kernal_1 = (np.array([0,1,0,0,0,0,0,1,0])/(16.0 * geo_trans[5] * eita)).reshape((3,3))
	kernal_2 = (np.array([0,0,0,1,0,1,0,0,0])/(16.0 * geo_trans[1] * eita)).reshape((3,3))
	kernal_3 = (np.array([1,0,1,0,0,0,1,0,1])/(2 * math.sqrt(2) * math.sqrt((eita * geo_trans[5]) ** 2 + (eita * geo_trans[1]) **2))).reshape((3,3))
	kernal_4 = (np.array([0,0,0,0,(math.sqrt(2)+1)/(8*math.sqrt(2)*math.sqrt((eita * geo_trans[5]) ** 2 + (eita * geo_trans[1]) **2)),0,0,0,0])).reshape((3,3))

	agg_1_dem = cv2.filter2D(dem_array,-1,kernal_1,borderType = 0)
	agg_2_dem = cv2.filter2D(dem_array,-1,kernal_2,borderType = 0)
	agg_3_dem = cv2.filter2D(dem_array,-1,kernal_3,borderType = 0)
	agg_4_dem = cv2.filter2D(dem_array,-1,kernal_4,borderType = 0)

	agg_5_dem = agg_4_dem - agg_1_dem + agg_2_dem + agg_3_dem
	# the range of curvature was scaled to [-0.5,0.5]
	return (agg_5_dem - (np.max(agg_5_dem)+np.min(agg_5_dem))/2.0)/(np.max(agg_5_dem) - np.min(agg_5_dem))
	
	
if __name__ == '__main__':
	# rfe_parse("/mnt/fire/Metero/rfe/cpc_rfe_v2.0_sa_dly.bin.20130102.gz")
	downscale_metero_filed("/mnt/fire/Metero","/mnt/fire/Extract","2012-12-01","2015-12-31")
	# raster.resample("/mnt/fire/Metero/geopotential/geopoptential.tif","/mnt/fire/Metero/dem/dem.tif","bilinear")
